"""
ROS Noetic node (Python) that integrates rs2/GStreamerWrapper video source
with the cone-alignment-filtered YOLO detection pipeline.

Behavior:
- Subscribes to /yolo_detect (std_msgs/String). When receives "start",
  it starts grabbing frames from GStreamer and running detection.
- For frames where cone, number, dashboard are all non-None, the result is
  added to a cache. When the cache reaches size 5, the majority vote across
  each category is computed and published to /yolo_result as
  "<cone>-<number>-<dashboard>", e.g., "red-five-up".
- After publishing, detection stops and waits for the next "start" command,
  and the cache is cleared.
"""

import os
import sys
import time
import threading
from collections import Counter
from typing import Optional, Dict, Any, Tuple, List

import numpy as np

# ROS imports
import rospy
from std_msgs.msg import String

# Ensure parent directory (yolo-jetson) is on sys.path to import rs2/*
CURRENT_DIR = os.path.dirname(__file__)
PARENT_DIR = os.path.abspath(os.path.join(CURRENT_DIR, '..'))
if PARENT_DIR not in sys.path:
    sys.path.append(PARENT_DIR)

from yolo_pipeline_optimized_cone_filtered import YOLOTrafficLightPipelineOptimized
from rs2.GStreamerWrapper import GStreamerWrapper


def _find_model_path() -> Optional[str]:
    candidate_models = [
        # Prefer TensorRT engine models
        r".\weights\all_yolo_all_epochs_0811\best.engine",
        r".\weights\all_yolo_all_epochs_0808\best.engine",
        r".\weights\all_yolo_all_epochs_0805\best.engine",
        r".\weights\all_yolo_all_epochs_0804\best.engine",
        # Fallback to PyTorch models
        r".\weights\all_yolo_all_epochs_0811\best.pt",
        r".\weights\all_yolo_all_epochs_0808\best.pt",
        r".\weights\all_yolo_all_epochs_0805\best.pt",
        r".\weights\all_yolo_all_epochs_0804\best.pt",
    ]
    for path in candidate_models:
        if os.path.exists(path):
            return path
    return None


class YoloGstConeFilteredNode:
    def __init__(self) -> None:
        # Params
        self.cache_size: int = int(rospy.get_param('~cache_size', 5))
        self.gst_width: int = int(rospy.get_param('~gst_width', 1280))
        self.gst_height: int = int(rospy.get_param('~gst_height', 720))
        self.device: Optional[str] = rospy.get_param('~device', None)
        self.use_fresh_pipeline_per_frame: bool = bool(rospy.get_param('~fresh_per_frame', True))

        # Model path
        model_path_param = rospy.get_param('~model_path', '')
        self.model_path: Optional[str] = model_path_param if model_path_param else _find_model_path()
        if not self.model_path or not os.path.exists(self.model_path):
            rospy.logerr('Model weights not found. Please set ~model_path or place weights under ./weights/.')
            raise RuntimeError('Model weights not found')

        # Preload once to validate model
        test_pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)
        del test_pipeline
        rospy.loginfo(f'Optimized model loaded: {self.model_path}')

        # ROS pub/sub
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.cmd_sub = rospy.Subscriber('/yolo_detect', String, self._cmd_callback, queue_size=10)

        # State
        self.active_lock = threading.Lock()
        self.active: bool = False
        self.stop_event = threading.Event()
        self.worker_thread: Optional[threading.Thread] = None
        self.result_cache: List[Tuple[str, str, str]] = []

        # Optional persistent pipeline
        self.persistent_pipeline: Optional[YOLOTrafficLightPipelineOptimized] = None
        if not self.use_fresh_pipeline_per_frame:
            self.persistent_pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)

    def _cmd_callback(self, msg: String) -> None:
        data = (msg.data or '').strip().lower()
        if data == 'start':
            with self.active_lock:
                if self.active:
                    rospy.loginfo('Start requested but detection is already active. Ignoring.')
                    return
                self.active = True
                self.stop_event.clear()
                self.result_cache.clear()
                self.worker_thread = threading.Thread(target=self._run_detection_loop, daemon=True)
                self.worker_thread.start()
                rospy.loginfo('Detection started.')
        elif data == 'stop':
            with self.active_lock:
                if self.active:
                    rospy.loginfo('Stop requested. Stopping detection...')
                    self.stop_event.set()
                else:
                    rospy.loginfo('Stop requested but detection is not active.')

    def _process_single_frame(self, frame: np.ndarray, frame_idx: int) -> Dict[str, Any]:
        try:
            if self.use_fresh_pipeline_per_frame:
                pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)
            else:
                pipeline = self.persistent_pipeline

            results = pipeline.process_image(frame)

            if self.use_fresh_pipeline_per_frame:
                del pipeline

            return results
        except Exception as e:
            rospy.logwarn(f'Error processing frame {frame_idx}: {e}')
            return {
                'final_results': [],
                'cone_color': 'None',
                'number': 'None',
                'dashboard_state': 'None',
                'error': str(e),
                'processing_info': {},
            }

    def _majority_vote(self, items: List[str]) -> str:
        if not items:
            return 'None'
        c = Counter(items)
        return c.most_common(1)[0][0]

    def _publish_and_stop(self) -> None:
        # Compute majority per field
        cones = [c for c, _, _ in self.result_cache]
        numbers = [n for _, n, _ in self.result_cache]
        dashboards = [d for _, _, d in self.result_cache]

        cone_final = self._majority_vote(cones)
        number_final = self._majority_vote(numbers)
        dashboard_final = self._majority_vote(dashboards)

        result_str = f'{cone_final}-{number_final}-{dashboard_final}'
        rospy.loginfo(f'Publishing result: {result_str}')
        self.result_pub.publish(String(data=result_str))

        # Stop and reset for next start
        self.stop_event.set()
        self.result_cache.clear()

    def _run_detection_loop(self) -> None:
        gst = None
        try:
            rospy.loginfo('Initializing GStreamer source...')
            gst = GStreamerWrapper(width=self.gst_width, height=self.gst_height)
            frame_idx = 0

            while not rospy.is_shutdown() and not self.stop_event.is_set():
                frame = gst.GetFrame()
                if frame is None:
                    time.sleep(0.005)
                    continue

                results = self._process_single_frame(frame, frame_idx)

                cone_color = results.get('cone_color', 'None')
                number = results.get('number', 'None')
                dashboard_state = results.get('dashboard_state', 'None')

                rospy.loginfo(f'Frame {frame_idx}: cone={cone_color}, number={number}, dashboard={dashboard_state}')

                if cone_color != 'None' and number != 'None' and dashboard_state != 'None':
                    self.result_cache.append((str(cone_color), str(number), str(dashboard_state)))
                    rospy.loginfo(f'Cached valid result {len(self.result_cache)}/{self.cache_size}')
                    if len(self.result_cache) >= self.cache_size:
                        self._publish_and_stop()
                        break

                frame_idx += 1
        except Exception as e:
            rospy.logerr(f'Detection loop error: {e}')
        finally:
            try:
                if gst is not None:
                    gst.StopThread()
            except Exception:
                pass
            with self.active_lock:
                self.active = False
            rospy.loginfo('Detection stopped. Waiting for next start.')


def main() -> None:
    rospy.init_node('yolo_gst_cone_filtered_node', anonymous=False)
    node = YoloGstConeFilteredNode()
    rospy.loginfo('yolo_gst_cone_filtered_node ready. Send "start" on /yolo_detect to begin. Optional: send "stop" to cancel.')
    rospy.spin()


if __name__ == '__main__':
    main()


